package io.kiki.stack.netty.channelManager;

import io.kiki.stack.netty.Constants;
import io.kiki.stack.netty.config.ConfigManager;
import io.kiki.stack.netty.exception.RemotingException;
import io.kiki.stack.netty.lifeCycle.AbstractLifeCycle;
import io.kiki.stack.netty.lifeCycle.LifeCycle;
import io.kiki.stack.netty.lifeCycle.LifeCycleException;
import io.kiki.stack.netty.util.AddressParser;
import io.kiki.stack.netty.util.NamedThreadFactory;
import io.kiki.stack.netty.util.StringUtils;
import lombok.Getter;
import lombok.NonNull;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.concurrent.*;

@Slf4j
@Setter
@Getter
public class ChannelManagerImpl extends AbstractLifeCycle implements ChannelManager, LifeCycle {


    /**
     * Channel pool initialize tasks
     */
    public ConcurrentMap<String, RunStateRecordedFutureTask<Pool>> poolKeyToRunStateRecordedFutureTaskMap;
    /**
     * heal Channel tasks
     */
    public ConcurrentMap<String, FutureTask<Integer>> healTasks;

    public ChannelSelectStrategy channelSelectStrategy;


    public ChannelFactory channelFactory;

    public ChannelEventHandler channelEventHandler;

    public ChannelEventListener channelEventListener;

    private ThreadPoolExecutor asyncCreateChannelExecutor;

    public ChannelManagerImpl(ChannelSelectStrategy channelSelectStrategy) {
        this.poolKeyToRunStateRecordedFutureTaskMap = new ConcurrentHashMap<String, RunStateRecordedFutureTask<Pool>>();
        this.healTasks = new ConcurrentHashMap<String, FutureTask<Integer>>();
        this.channelSelectStrategy = channelSelectStrategy;
    }

    public ChannelManagerImpl(ChannelSelectStrategy channelSelectStrategy, ChannelFactory channelFactory) {
        this(channelSelectStrategy);
        this.channelFactory = channelFactory;
    }

    public ChannelManagerImpl(ChannelFactory channelFactory, ChannelEventHandler channelEventHandler) {
        this(new RandomChannelSelectStrategy(null), channelFactory);

        this.channelEventHandler = channelEventHandler;
    }

    public ChannelManagerImpl(ChannelSelectStrategy channelSelectStrategy, ChannelFactory channelFactory, ChannelEventHandler channelEventHandler, ChannelEventListener channelEventListener) {
        this(channelSelectStrategy, channelFactory);
        this.channelEventHandler = channelEventHandler;
        this.channelEventListener = channelEventListener;
    }

    @Override
    public void startup() throws LifeCycleException {
        super.startup();

        long keepAliveTime = ConfigManager.conn_create_tp_keepalive();
        int queueSize = ConfigManager.conn_create_tp_queue_size();
        int minPoolSize = ConfigManager.conn_create_tp_min_size();
        int maxPoolSize = ConfigManager.conn_create_tp_max_size();
        this.asyncCreateChannelExecutor = new ThreadPoolExecutor(minPoolSize, maxPoolSize, keepAliveTime, TimeUnit.SECONDS, new ArrayBlockingQueue<Runnable>(queueSize), new NamedThreadFactory("-conn-warmup-executor", true));
    }

    @Override
    public void shutdown() throws LifeCycleException {
        super.shutdown();

        if (asyncCreateChannelExecutor != null) {
            asyncCreateChannelExecutor.shutdown();
        }

        if (null == this.poolKeyToRunStateRecordedFutureTaskMap || this.poolKeyToRunStateRecordedFutureTaskMap.isEmpty()) {
            return;
        }
        Iterator<String> iterator = this.poolKeyToRunStateRecordedFutureTaskMap.keySet().iterator();
        while (iterator.hasNext()) {
            String poolKey = iterator.next();
            this.removeTask(poolKey);
            iterator.remove();
        }
        log.warn("All Channel pool and Channels have been removed!");
    }

    @Deprecated
    @Override
    public void init() {
        this.channelEventHandler.setChannelManager(this);
        this.channelEventHandler.setChannelEventListener(channelEventListener);
        this.channelFactory.init(channelEventHandler);
    }

    @Override
    public void add(Channel channel) {
        Set<String> poolKeys = channel.getPoolKeys();
        for (String poolKey : poolKeys) {
            this.add(channel, poolKey);
        }
    }

    @Override
    public void add(Channel channel, String poolKey) {
        Pool pool = null;
        try {
            // get or create an empty Channel pool
            pool = this.getChannelPoolAndCreateIfAbsent(poolKey, new CreatePoolTask(this));
        } catch (Exception e) {
            // should not reach here.
            log.error("[NOTIFYME] Exception occurred when getOrCreateIfAbsent an empty channelPool!", e);
        }
        if (pool != null) {
            pool.add(channel);
        } else {
            // should not reach here.
            log.error("[NOTIFYME] Channel pool NULL!");
        }
    }

    @Override
    public Channel get(String poolKey) {
        Pool pool = this.getPool(this.poolKeyToRunStateRecordedFutureTaskMap.get(poolKey));
        return null == pool ? null : pool.get();
    }

    @Override
    public List<Channel> getAll(String poolKey) {
        Pool pool = this.getPool(this.poolKeyToRunStateRecordedFutureTaskMap.get(poolKey));
        return null == pool ? new ArrayList<Channel>() : pool.getAll();
    }

    @Override
    public Map<String, List<Channel>> getAll() {
        Map<String, List<Channel>> map = new HashMap<String, List<Channel>>();
        for (Map.Entry<String, RunStateRecordedFutureTask<Pool>> entry : this.getPoolKeyToRunStateRecordedFutureTaskMap().entrySet()) {
            Pool pool = RunStateRecordedFutureTask.getFutureTaskResult(entry.getValue());
            if (null != pool) {
                map.put(entry.getKey(), pool.getAll());
            }
        }
        return map;
    }

    @Override
    public void remove(Channel channel) {
        if (null == channel) {
            return;
        }
        Set<String> poolKeys = channel.getPoolKeys();
        if (null == poolKeys || poolKeys.isEmpty()) {
            channel.close();
            log.warn("Remove and close a standalone Channel.");
        } else {
            for (String poolKey : poolKeys) {
                this.remove(channel, poolKey);
            }
        }
    }

    @Override
    public void remove(Channel channel, String poolKey) {
        if (null == channel || StringUtils.isBlank(poolKey)) {
            return;
        }
        Pool pool = this.getPool(this.poolKeyToRunStateRecordedFutureTaskMap.get(poolKey));
        if (null == pool) {
            channel.close();
            log.warn("Remove and close a standalone Channel.");
        } else {
            pool.removeAndTryClose(channel);
            if (pool.isEmpty()) {
                this.removeTask(poolKey);
                log.warn("Remove and close the last Channel in channelPool with poolKey {}", poolKey);
            } else {
                log.warn("Remove and close a channel in channelPool with poolKey {}, {} channels left.", poolKey, pool.size());
            }
        }
    }

    @Override
    public void remove(String poolKey) {
        if (StringUtils.isBlank(poolKey)) {
            return;
        }

        RunStateRecordedFutureTask<Pool> task = this.poolKeyToRunStateRecordedFutureTaskMap.remove(poolKey);
        if (null != task) {
            Pool pool = this.getPool(task);
            if (null != pool) {
                pool.removeAllAndTryClose();
                log.warn("Remove and close all channels in channelPool of poolKey={}", poolKey);
            }
        }
    }

    /**
     * Warning! This is weakly consistent implementation, to prevent lock the whole {@link ConcurrentHashMap}.
     */
    @Deprecated
    @Override
    public void removeAll() {
        if (null == this.poolKeyToRunStateRecordedFutureTaskMap || this.poolKeyToRunStateRecordedFutureTaskMap.isEmpty()) {
            return;
        }
        Iterator<String> iterator = this.poolKeyToRunStateRecordedFutureTaskMap.keySet().iterator();
        while (iterator.hasNext()) {
            String poolKey = iterator.next();
            this.removeTask(poolKey);
            iterator.remove();
        }
        log.warn("All channel pool and channels have been removed!");
    }

    @Override
    public void check(@NonNull Channel channel) throws RemotingException {
        io.netty.channel.Channel channel1 = channel.getChannel();
        if (channel1 == null || !channel1.isActive()) {
            this.remove(channel);
            throw new RemotingException("Check channel failed for address: " + channel.getUrl());
        }
        if (!channel1.isWritable()) {
            // No remove. Most of the time it is unwritable temporarily.
            throw new RemotingException("Check channel failed for address: " + channel.getUrl() + ", maybe write overflow!");
        }
    }

    @Override
    public int count(String poolKey) {
        if (StringUtils.isBlank(poolKey)) {
            return 0;
        }
        Pool pool = this.getPool(this.poolKeyToRunStateRecordedFutureTaskMap.get(poolKey));
        if (null != pool) {
            return pool.size();
        } else {
            return 0;
        }
    }

    /**
     * in case of cache pollution and channel leak, to do schedule scan
     */
    @Override
    public void scan() {
        if (null != this.poolKeyToRunStateRecordedFutureTaskMap && !this.poolKeyToRunStateRecordedFutureTaskMap.isEmpty()) {
            Iterator<String> iterator = this.poolKeyToRunStateRecordedFutureTaskMap.keySet().iterator();
            while (iterator.hasNext()) {
                String poolKey = iterator.next();
                RunStateRecordedFutureTask<Pool> task = this.poolKeyToRunStateRecordedFutureTaskMap.get(poolKey);
                if (task == null) {
                    log.warn("task(poolKey={}) is null, do not scan the channel pool", poolKey);
                    continue;
                }

                if (!task.isDone()) {
                    log.info("task(poolKey={}) is not done, do not scan the channel pool", poolKey);
                    continue;
                }

                Pool pool = this.getPool(task);
                if (null != pool) {
                    pool.scan();
                    if (pool.isEmpty()) {
                        if ((System.currentTimeMillis() - pool.getLastAccessTimestamp()) > Constants.default_expire_time) {
                            iterator.remove();
                            log.warn("Remove expired pool task of poolKey {} which is empty.", poolKey);
                        }
                    }
                }
            }
        }
    }

    /**
     * If no task cached, create one and initialize the channels.
     */
    @Override
    public Channel getAndCreateIfAbsent(Url url) throws InterruptedException, RemotingException {
        // get and create a channel pool with initialized channels.
        Pool pool = this.getChannelPoolAndCreateIfAbsent(url.getId(), new CreatePoolTask(this, url));
        if (null != pool) {
            return pool.get();
        } else {
            log.error("[NOTIFYME] bug detected! pool here must not be null!");
            return null;
        }
    }

    @Override
    public void asyncCreateChannel(Url url) throws RemotingException, InterruptedException {
        this.getChannelPoolAndCreateIfAbsent(url.getId(), new CreatePoolTask(this, url, 0));
    }

    /**
     * If no task cached, create one and initialize the channels.
     * If task cached, check whether the number of channels adequate, if not then heal it.
     */
    @Override
    public void createChannelAndHealIfNeed(Url url) throws InterruptedException, RemotingException {
        // get and create a channel pool with initialized channels.
        Pool pool = this.getChannelPoolAndCreateIfAbsent(url.getId(), new CreatePoolTask(this, url));
        if (null != pool) {
            healPoolIfNeed(pool, url);
        } else {
            log.error("[NOTIFYME] bug detected! pool here must not be null!");
        }
    }

    @Override
    public Channel create(Url url) throws RemotingException {
        Channel channel;
        try {
            channel = this.channelFactory.createChannel(url);
        } catch (Exception e) {
            throw new RemotingException("Create channel failed. The address is " + url.getOriginUrl(), e);
        }
        return channel;
    }

    @Override
    public Channel create(String ip, int port, int connectTimeout) throws RemotingException {
        try {
            return this.channelFactory.createChannel(ip, port, connectTimeout);
        } catch (Exception e) {
            throw new RemotingException("Create channel failed. The address is " + ip + ":" + port, e);
        }
    }

    @Override
    public Channel create(String address, int connectTimeout) throws RemotingException {
        Url url = AddressParser.parse(address);
        url.setConnectTimeout(connectTimeout);
        return create(url);
    }

    @Override
    public void disableHeartbeat(Channel channel) {
        if (null != channel) {
            channel.getChannel().attr(Channel.HEARTBEAT_SWITCH).set(false);
        }
    }

    @Override
    public void enableHeartbeat(Channel channel) {
        if (null != channel) {
            channel.getChannel().attr(Channel.HEARTBEAT_SWITCH).set(true);
        }
    }

    private Pool getPool(RunStateRecordedFutureTask<Pool> task) {
        return RunStateRecordedFutureTask.getFutureTaskResult(task);

    }

    private Pool getChannelPoolAndCreateIfAbsent(String poolKey, Callable<Pool> callable) throws RemotingException, InterruptedException {
        RunStateRecordedFutureTask<Pool> task;
        Pool pool = null;

        int retry = Constants.default_retry_times;

        int timesOfResultNull = 0;
        int timesOfInterrupt = 0;

        for (int i = 0; (i < retry) && (pool == null); ++i) {
            task = this.poolKeyToRunStateRecordedFutureTaskMap.get(poolKey);
            if (null == task) {
                RunStateRecordedFutureTask<Pool> newTask = new RunStateRecordedFutureTask<Pool>(callable);
                task = this.poolKeyToRunStateRecordedFutureTaskMap.putIfAbsent(poolKey, newTask);
                if (null == task) {
                    task = newTask;
                    task.run();
                }
            }

            try {
                pool = task.get();
                if (null == pool) {
                    if (i + 1 < retry) {
                        timesOfResultNull++;
                        continue;
                    }
                    this.poolKeyToRunStateRecordedFutureTaskMap.remove(poolKey);
                    String errMsg = "Get future task result null for poolKey [" + poolKey + "] after [" + (timesOfResultNull + 1) + "] times try.";
                    throw new RemotingException(errMsg);
                }
            } catch (InterruptedException e) {
                if (i + 1 < retry) {
                    timesOfInterrupt++;
                    continue;// retry if interrupted
                }
                this.poolKeyToRunStateRecordedFutureTaskMap.remove(poolKey);
                log.warn("Future task of poolKey {} interrupted {} times. InterruptedException thrown and stop retry.", poolKey, (timesOfInterrupt + 1), e);
                throw e;
            } catch (ExecutionException e) {
                // DO NOT retry if ExecutionException occurred
                this.poolKeyToRunStateRecordedFutureTaskMap.remove(poolKey);

                Throwable cause = e.getCause();
                if (cause instanceof RemotingException) {
                    throw (RemotingException) cause;
                } else {
                    RunStateRecordedFutureTask.launderThrowable(cause);
                }
            }
        }
        return pool;
    }

    /**
     * remove task and remove all channels
     *
     * @param poolKey target pool key
     */
    protected void removeTask(String poolKey) {
        RunStateRecordedFutureTask<Pool> task = this.poolKeyToRunStateRecordedFutureTaskMap.remove(poolKey);
        if (null != task) {
            Pool pool = RunStateRecordedFutureTask.getFutureTaskResult(task);
            if (null != pool) {
                pool.removeAllAndTryClose();
            }
        }
    }

    /**
     * execute heal channel tasks if the actual number of channels in pool is less than expected
     *
     * @param pool channel pool
     * @param url  target url
     */
    private void healPoolIfNeed(Pool pool, Url url) throws RemotingException, InterruptedException {
        String poolKey = url.getId();
        // only when async creating channels done
        // and the actual size of channels less than expected, the healing task can be run.
        if (pool.isAsyncCreationDone() && pool.size() < url.getChannelNum()) {
            FutureTask<Integer> task = this.healTasks.get(poolKey);
            if (null == task) {
                FutureTask<Integer> futureTask = new FutureTask<Integer>(new HealPoolTask(this, url, pool));
                task = this.healTasks.putIfAbsent(poolKey, futureTask);
                if (null == task) {
                    task = futureTask;
                    task.run();
                }
            }
            try {
                int numAfterHeal = task.get();
                if (log.isDebugEnabled()) {
                    log.debug("[NOTIFYME] - conn num after heal {}, expected {}, warmup {}", numAfterHeal, url.getChannelNum(), url.isConnWarmup());
                }
            } catch (InterruptedException e) {
                this.healTasks.remove(poolKey);
                throw e;
            } catch (ExecutionException e) {
                this.healTasks.remove(poolKey);
                Throwable cause = e.getCause();
                if (cause instanceof RemotingException) {
                    throw (RemotingException) cause;
                } else {
                    RunStateRecordedFutureTask.launderThrowable(cause);
                }
            }
            // heal task is one-off, remove from cache directly after run
            this.healTasks.remove(poolKey);
        }
    }

    /**
     * @param syncCreateNumWhenNotWarmup you can specify this param to ensure at least desired number of channels available in sync way
     */
    public void doCreatePool(final Pool pool, final Url url, final int syncCreateNumWhenNotWarmup) throws RemotingException {
        final int actualNum = pool.size();
        final int expectNum = url.getChannelNum();
        if (actualNum >= expectNum) {
            if (log.isDebugEnabled()) {
                log.debug("actual num {}, expect num {}", actualNum, expectNum);
            }
            return;
        }
        if (url.isConnWarmup()) {
            for (int i = actualNum; i < expectNum; ++i) {
                Channel channel = create(url);
                pool.add(channel);
            }
        } else {
            if (syncCreateNumWhenNotWarmup < 0 || syncCreateNumWhenNotWarmup > url.getChannelNum()) {
                throw new IllegalArgumentException("sync create number when not warmup should be [0," + url.getChannelNum() + "]");
            }
            // create channel in sync way
            if (syncCreateNumWhenNotWarmup > 0) {
                for (int i = 0; i < syncCreateNumWhenNotWarmup; ++i) {
                    Channel channel = create(url);
                    pool.add(channel);
                }
                if (syncCreateNumWhenNotWarmup >= url.getChannelNum()) {
                    return;
                }
            }

            pool.markAsyncCreationStart();// mark the start of async
            try {
                this.asyncCreateChannelExecutor.execute(new Runnable() {
                    @Override
                    public void run() {
                        try {
                            for (int i = pool.size(); i < url.getChannelNum(); ++i) {
                                try {
                                    Channel channel = create(url);
                                    pool.add(channel);
                                } catch (RemotingException e) {
                                    log.error("Exception occurred in async create channel thread for {}, taskName {}", url.getId(), e);
                                    if (pool.isEmpty()) {
                                        poolKeyToRunStateRecordedFutureTaskMap.remove(url.getId());
                                        break;
                                    }
                                }
                            }
                        } finally {
                            pool.markAsyncCreationDone();// mark the end of async
                        }
                    }
                });
            } catch (RejectedExecutionException e) {
                pool.markAsyncCreationDone();// mark the end of async when reject
                throw e;
            }
        } // end of NOT warm up
    }


}
